import {
    ChatAgent,
    ChatAgentConfig,
    RawChatAgentConfig,
    ChatAgentResponse
} from "../agent/chatAgent";
import { TaskPlannerAgent, TaskSpecifyAgent } from "../agent/taskAgent";
import { Language, RoleType } from "../enum";
import { ChatMessage, ReceivingChatMessage } from "../message/chatMessage";
import { SystemMessage } from "../message/systemMessage";
import { personConfigs } from "../config";
import { ModelBackendConfig } from "../model";
import { talk } from "./agentChat";

// 定义 RolePlayingConfig 接口
interface RolePlayingConfig {
    /* 用户角色名称 */
    userRoleName: string;
    /* 助手角色名称 */
    assistantRoleName: string;
    /* 用户角色类型，默认为 null */
    userRoleType?: RoleType;
    /* 助手角色类型，默认为 null */
    assistantRoleType?: RoleType;
    /* 任务提示，默认为空字符串 */
    taskPrompt?: string;
    /* 背景提示，默认为 null */
    backgroundPrompt?: string;
    /* 用户角色提示，默认为空字符串 */
    userRolePrompt?: string;
    /* 助手角色提示，默认为空字符串 */
    assistantRolePrompt?: string;
    /* 是否使用任务细化，默认为 true */
    withTaskSpecify?: boolean;
    /* 是否使用任务规划，默认为 false */
    withTaskPlanner?: boolean;
    /* 模型配置 */
    modelConfig?: ModelBackendConfig;
    /* 传递给助手代理的额外参数，默认为 null */
    assistantAgentConfig?: ChatAgentConfig;
    /* 传递给用户代理的额外参数，默认为 null */
    userAgentConfig?: ChatAgentConfig;
    /* 传递给任务细化代理的额外参数，默认为 null */
    taskSpecifyAgentConfig?: RawChatAgentConfig;
    /* 传递给任务规划代理的额外参数，默认为 null */
    taskPlannerAgentConfig?: RawChatAgentConfig;
    /* 公共记忆，默认为 null */
    memory?: any;
    /* 记忆深度，默认为 1 */
    memoryDepth?: number;
    /* 语言，默认 Chinese */
    language?: Language;
}

enum RolePlayingAgent {
    User = "user",
    Assistant = "assistant",
    TaskPlanner = "taskPlanner",
    TaskSpecify = "taskSpecify"
}

// 定义 RolePlaying 类
export class RolePlaying {
    public userRoleName: string;
    public assistantRoleName: string;
    public userRoleType?: RoleType;
    public assistantRoleType?: RoleType;
    public taskPrompt?: string;
    public backgroundPrompt?: string;
    public userRolePrompt?: string;
    public assistantRolePrompt?: string;
    public withTaskSpecify?: boolean;
    public withTaskPlanner?: boolean;
    public modelConfig?: ModelBackendConfig;
    public assistantAgentConfig?: ChatAgentConfig;
    public userAgentConfig?: ChatAgentConfig;
    public taskSpecifyAgentConfig?: RawChatAgentConfig;
    public taskPlannerAgentConfig?: RawChatAgentConfig;
    public memory?: any;
    public memoryDepth?: number;
    public language?: Language;

    /** Agent 字典 */
    private agentDict = {
        [RolePlayingAgent.User]: ChatAgent,
        [RolePlayingAgent.Assistant]: ChatAgent,
        [RolePlayingAgent.TaskSpecify]: TaskSpecifyAgent,
        [RolePlayingAgent.TaskPlanner]: TaskPlannerAgent
    };
    /** Agent 缓存字典 */
    private agentTempDict: {
        [RolePlayingAgent.User]?: ChatAgent;
        [RolePlayingAgent.Assistant]?: ChatAgent;
        [RolePlayingAgent.TaskSpecify]?: TaskSpecifyAgent;
        [RolePlayingAgent.TaskPlanner]?: TaskPlannerAgent;
    } = {};

    constructor(config: RolePlayingConfig) {
        this.userRoleName = config.userRoleName;
        this.assistantRoleName = config.assistantRoleName;
        this.userRoleType = config.userRoleType ?? RoleType.User;
        this.assistantRoleType = config.assistantRoleType ?? RoleType.Assistant;
        this.taskPrompt = config.taskPrompt;
        this.backgroundPrompt = config.backgroundPrompt;
        this.userRolePrompt = config.userRolePrompt ?? "";
        this.assistantRolePrompt = config.assistantRolePrompt ?? "";
        this.withTaskSpecify = config.withTaskSpecify ?? true;
        this.withTaskPlanner = config.withTaskPlanner ?? false;
        this.modelConfig = config.modelConfig ?? personConfigs.defaultModelBackendConfig;
        this.assistantAgentConfig = config.assistantAgentConfig;
        this.userAgentConfig = config.userAgentConfig;
        this.taskSpecifyAgentConfig = config.taskSpecifyAgentConfig;
        this.taskPlannerAgentConfig = config.taskPlannerAgentConfig;
        this.memory = config.memory ?? null;
        this.memoryDepth = config.memoryDepth ?? 1;
        this.language = config.language ?? personConfigs.defaultLanguage;
    }

    /**
     * 初始化 Agent
     */
    async initChat(): Promise<{
        [RolePlayingAgent.User]: ChatAgent;
        [RolePlayingAgent.Assistant]: ChatAgent;
    }> {
        const userAgent = this.agentTempDict[RolePlayingAgent.User];
        const assistantAgent = this.agentTempDict[RolePlayingAgent.Assistant];
        if (userAgent && assistantAgent) {
            userAgent.reset();
            assistantAgent.reset();
        } else {
            this.agentTempDict[RolePlayingAgent.Assistant] = new this.agentDict[
                RolePlayingAgent.Assistant
            ]({
                ...this.assistantAgentConfig,
                systemMessage: new SystemMessage({
                    roleName: this.assistantRoleName,
                    roleType: this.assistantRoleType ?? RoleType.Assistant,
                    content: this.assistantRolePrompt ?? ""
                }),
                language: this.language,
                memory: this.memory
            });
            this.agentTempDict[RolePlayingAgent.User] = new this.agentDict[RolePlayingAgent.User]({
                ...this.userAgentConfig,
                systemMessage: new SystemMessage({
                    roleName: this.userRoleName,
                    roleType: this.userRoleType ?? RoleType.User,
                    content: this.userRolePrompt ?? ""
                }),
                language: this.language,
                memory: this.memory
            });
        }
        return {
            [RolePlayingAgent.User]: this.agentTempDict[RolePlayingAgent.User]!,
            [RolePlayingAgent.Assistant]: this.agentTempDict[RolePlayingAgent.Assistant]!
        };
    }
    /**
     * 执行一步对话
     * @param inputMessage 输入消息
     * @returns 所有 Agent 的回复
     */
    async step(inputMessage: ChatMessage): Promise<{
        [RolePlayingAgent.User]: ChatAgentResponse;
        [RolePlayingAgent.Assistant]: ChatAgentResponse;
    }> {
        const { [RolePlayingAgent.User]: initiator, [RolePlayingAgent.Assistant]: responder } =
            await this.initChat();
        // 优化提示
        inputMessage.content = await this.optimizePrompt(inputMessage.content);
        // 交谈
        const { initiatorRes, responderRes } = await talk(
            initiator,
            responder,
            inputMessage,
            this.memory ?? true,
            this.memoryDepth
        );
        if (!initiatorRes || !responderRes) {
            throw new Error("Chat agent response is empty.");
        }
        return {
            [RolePlayingAgent.User]: initiatorRes,
            [RolePlayingAgent.Assistant]: responderRes
        };
    }

    private async useTaskAgent(
        rolePlayingAgent: RolePlayingAgent,
        agentConfig: any,
        taskPrompt: string,
        backgroundPrompt?: string
    ): Promise<string> {
        if (!this.agentTempDict[rolePlayingAgent]) {
            this.agentTempDict[rolePlayingAgent] = new this.agentDict[rolePlayingAgent]({
                ...agentConfig
            });
        }
        const agent = this.agentTempDict[rolePlayingAgent];

        const prompt = `${taskPrompt}\n${backgroundPrompt ?? ""}`;
        const result = await agent.step(
            new ReceivingChatMessage({
                roleName: "RolePlaying",
                roleType: RoleType.Assistant,
                content: prompt
            })
        );

        if (result.terminated) {
            return prompt;
        } else {
            return agent.getLatestMessage().content;
        }
    }
    private async optimizePrompt(prompt: string): Promise<string> {
        let promptResult = prompt;
        if (this.withTaskSpecify && this.taskPrompt) {
            promptResult = await this.useTaskAgent(
                RolePlayingAgent.TaskSpecify,
                this.taskSpecifyAgentConfig,
                promptResult,
                this.backgroundPrompt
            );
        }
        if (this.withTaskPlanner && this.taskPrompt) {
            promptResult = await this.useTaskAgent(
                RolePlayingAgent.TaskPlanner,
                this.taskPlannerAgentConfig,
                promptResult,
                this.backgroundPrompt
            );
        }
        return promptResult;
    }
}
